// Copyright (C) 2010 and onwards Noam Shazeer and Georges Harik
// Pieces of this code may be licenced under the Apache License, version 2.0 
//   from Harik Shazeer Labs

// TODO: what is this needed for?
#ifndef _GNU_SOURCE 
#define _GNU_SOURCE
#endif

#include <unistd.h>
#include <sys/mman.h>
#include "allocator.h"
#include "jobqueue.h"
#include "interpret.h"

// Global objects
GlobalAllocator * G_ALLOC;
Sync * G_SYNC = NULL;

void InitAllocator() {
  if (G_ALLOC) return;
  G_SYNC = new Sync();
  G_SYNC->AddMaintenanceRoutine(&GarbageCollectMaintenance);
  InitJobQueue();
  G_ALLOC = new GlobalAllocator; 
  if (G_CONFIG->Exists("start_from")) {
    string prefix = G_CONFIG->GetString("start_from");
    G_ALLOC->Load(prefix + ".memory");
  }
}

GlobalAllocator::GlobalAllocator() {
  needs_garbage_collection_ = false;
  box_space_ = 1ull << ConfigInt("lg_box_space");
  emergency_box_space_ = 1ull << ConfigInt("lg_emergency_box_space");
  block_size_ = 1ull << ConfigInt("lg_global_allocation_block_size");
  uint64 space_size = box_space_ + emergency_box_space_;
  uint64 * space = new uint64[space_size];
  block_ = AllocationBlock::Make(space, space_size);
  G_BOX_INDEX_SPACE = 1ull << ConfigInt("lg_box_index_space");
  G_BOX_INDEX_MASK = G_BOX_INDEX_SPACE - 1;
  G_BOX_INDEX = new Box *[G_BOX_INDEX_SPACE];
  ClearBoxIndex();
  pthread_mutex_init(&lock_, NULL);
}

void GlobalAllocator::ClearBoxIndex() {
  // cout << "::ClearBoxIndex" << endl;
  JobQueueZero(G_JQ, G_BOX_INDEX, G_BOX_INDEX_SPACE * sizeof(Box *)); 
  cout << "ClearBoxIndex::" << endl;
}

// Returns NULL on failure
AllocationBlock * 
GlobalAllocator::AllocateBlock(uint64 length, bool force) {
  MutexLocker Locker(&lock_);
  if ((block_->used_ + length) > box_space_) {
    cout << "Not Enough Space in Global Allocator" << endl;
    needs_garbage_collection_ = true;
    if (!force) return NULL;
    cout << "Forced allocation of size " << length << endl;
  }
  AllocationBlock * ret = block_->AllocateSubblock(length);
  if (!ret) {
    cout << "Failure - legnth = " << length 
	 << " force = " << (force ? 'T' : 'F')
	 << " block_->used_ " << block_->used_
	 << " block_->capacity_ " << block_->capacity_ << endl;
    CHECK(false);
  }
  return ret;
}

ThreadAllocator::ThreadAllocator() 
  : Sync::Worker(G_SYNC), cache_(1 << ConfigInt("thread_cache_lg_size")) {
  global_block_ = NULL; local_block_ = NULL;
  startup_block_size_ = 64;
  uint64 local_size = 1 << ConfigInt("lg_local_allocation_block_size");
  ReallocateLocalBlock(local_size);
  global_words_allocated_ = 0;
  // generation_ = 1001;
  // local_block_->SetAllMemory(1000);
}

ThreadAllocator::~ThreadAllocator() {
  delete[] (uint64*) local_block_; 
}

uint64 * ThreadAllocator::GlobalAllocate(uint64 length) {
  CheckStop();
  global_words_allocated_ += length;
  if (global_block_) {
    uint64 * ret = global_block_->Allocate(length);
    if (ret) return ret;
  }
  uint64 block_size = min(startup_block_size_, G_ALLOC->block_size_);
  startup_block_size_ = 2 * block_size;
  uint64 size = max(length, block_size);
  global_block_ = G_ALLOC->AllocateBlock(size, critical_); 
  if (!global_block_) StopAll();
  uint64 * ret =  global_block_->Allocate(length);
  CHECK(ret);
  return ret;
}

uint64 * ThreadAllocator::LocalAllocate(uint64 length) {
  CheckStop();
  uint64 * ret = local_block_->Allocate(length);
  if (!ret) {
    throw LocalAllocatorException();
  }
  // for (uint64 i=0; i<length; i++) CHECK(ret[i] < generation_);
  return ret;
}

void ThreadAllocator::IndexBox(Box *b) {
  uint64 steps = b->Index();
  if (steps >= 50) {
    G_ALLOC->needs_garbage_collection_ = true;
    StopAll(); } }

Thing ThreadAllocator::MakeImmutableArray(AllocationMode m, 
					  Thing* contents, 
					  uint64 length) {
  uint64 id = Thing::MakeArrayID(contents, length);
  if (m == LOCAL_ALLOCATION) {
    uint64 * space = LocalAllocate(length + (sizeof(Box) / sizeof(uint64)));
    Box * b = (Box *)space;
    b->InitImmutableArray(id, contents, length);
    return (uint64)b; 
  }

  // Globally allocated array
  Thing thing = id;
  Box **bpp = thing.GetBoxPointerPointer();
  if (bpp) {
    Box *old_box = *bpp;
    if (!(old_box->IsEqualToImmutableArray(contents, length))) {
	cout << "::MakeImmutableArray *** CollisionException ***" << endl;
	throw CollisionException();
    }
    return thing;
  }
  uint64 * space = GlobalAllocate(length + (sizeof(Box) / sizeof(uint64)));
  Box * fresh_box = (Box *)space;
  CHECK(!ContainsDirectArrays(contents, length))
  fresh_box->InitImmutableArray(id, contents, length);
  IndexBox(fresh_box);
  return thing; 
}

Thing ThreadAllocator::MakeMutableArray() {
  uint64 * space = LocalAllocate(sizeof(Box) / sizeof(uint64));
  Box * b = (Box *) space;
  b->InitMutableArray();
  return (uint64)b; 
}

// Leaves resized unused region uninitialized
void ThreadAllocator::ResizeMutableArray(Box *a, uint64 new_length) {
  uint64 old_length = a->GetLength();
  if (old_length == new_length) return;
  a->SetLength(new_length);
  if (new_length < old_length) return;
  if (new_length > a->marray_capacity__) {
    uint64 new_capacity = max(new_length, old_length * 2);
    uint64 * new_space = LocalAllocate(new_capacity);
    memcpy(new_space, a->MutableArrayData(), old_length * sizeof(uint64));
    a->marray_capacity__ = new_capacity;
    a->marray_data__ = (Box *)new_space;
  }
}

Thing ThreadAllocator::MakeMutableArray(uint64 length, Thing initial_element) {
  Thing ret = MakeMutableArray();
  Box * retbox = ret.GetBox();
  ResizeMutableArray(retbox, length);
  for (uint64 i=0; i<length; i++) retbox->Ref(i) = initial_element;
  return ret;
}

Thing ThreadAllocator::MakeMutableArray(Thing *contents, uint64 length) {
  Thing ret = MakeMutableArray();
  Box * retbox = ret.GetBox();
  ResizeMutableArray(retbox, length);
  for (uint64 i=0; i<length; i++) retbox->Ref(i) = contents[i];
  return ret;
}

Thing ThreadAllocator::MakeSymbol(const string & s) {
  Thing thing = Thing::MakeID(SYMBOL, Fingerprint(s));
  if (thing.MaybeGetBox()) return thing;
  uint64 * space = GlobalAllocate(s.size() + (sizeof(Box) / sizeof(uint64)));
  Box * b = (Box *) space;
  b->InitSymbol(thing.ID(), s);
  IndexBox(b);
  return thing; }

Thing ThreadAllocator::Globalize(Thing t) {
  // cout << "Globalize " << (void *)t.data_ << " " << t << endl;
  if (t.GetBasicType() != DIRECT_ARRAY) return t;
  CHECK(t.GetType() == IMMUTABLE_ARRAY);
  // head will be a head of a stack of boxes to globalize.
  // we pop a box off the stack after we are sure there is a global version. 
  Box * head = BoxListEnd();
  BoxListPush(&head, t.GetBox());
  while (head != BoxListEnd()) {
    // BoxListDisplay(head);
    Box * b = head;
    // cout << "Head = " << b << " " << b->LocalBoxToThing() << endl;
    Box * found_box = Thing(b->id_).MaybeGetBox();
    if (found_box) {
      // cout << "Found box " << endl;
      if (BoxesHaveSameIDs(b, found_box)) {
	BoxListPop(&head);
	continue;
      }
      // cout << "::Globalize *** CollisionException ***" << endl;
      throw CollisionException();
    }
    // cout << "Box not found " << endl;
    bool has_direct_array_child = false;
    Box * shove_after_child = NULL;
    for (uint64 i=0; i<b->ArraySize(); i++) {
      Thing child = b->Ref(i);
      if (child.GetBasicType() == DIRECT_ARRAY) {
	Box * found_child = Thing(child.ID()).MaybeGetBox();
	if (!found_child) {
	  // cout << "Enqueue child " << child << endl;
	  bool listpush = BoxListPush(&head, child.GetBox());
	  if (!listpush) {
	    // cout << "no listpush " << child << endl;
	    shove_after_child = child.GetBox();
	  }
	  has_direct_array_child = true;
	}
      }
    }
    if (shove_after_child && (head == b)) 
      BoxListShoveHeadAfter(&head, shove_after_child);
    if (!has_direct_array_child) {
      // cout << "Pop " << b->LocalBoxToThing() << endl;
      BoxListPop(&head);
      for (uint64 i=0; i<b->ArraySize(); i++) b->Ref(i) = b->Ref(i).ID();
      MakeImmutableArray(GLOBAL_ALLOCATION, 
			 b->ImmutableArrayData(), 
			 b->ArraySize());
    }
  }
  return t.ID();
}

// ---------- Garbage Collection from here on ----------

struct ExpandFrontierTask : public JobQueue::Task {
  ExpandFrontierTask(Frontier frontier) : frontier_(frontier) {}
  void Run() {
    // load balancing not yet implemented.
    CHECK(frontier_.start_ == frontier_.end_);
    // Add all things in a range to the linked list
    //frontier_.ExpandRange(frontier_.start_, frontier_.end_);
    while (frontier_.head_ != (Box *)0x1) {
      SplitWhileNecessary();
      frontier_.Expand();
    }
    __sync_fetch_and_add(&G_ALLOC->memory_used_, frontier_.seen_memory_);
    __sync_fetch_and_add(&G_ALLOC->boxes_used_, frontier_.seen_boxes_);
    // cout << "Frontier task done" << endl;
  }
  // TODO: For load balancing implement this
  void SplitWhileNecessary() {}
  Frontier frontier_;
};

// Insert reachable objects for a shard of an AllocatorDependant
// into a frontier and push it onto the global allocator
struct FlagReferencesTask : public JobQueue::Task {
  FlagReferencesTask (AllocatorDependant * ad, 
		      int32 shard)
    : ad_(ad), shard_(shard) {}
  void Run() { 
    ad_->FlagReferences(shard_, &frontier_);
    G_ALLOC->propagate_job_->AddTask
      (new ExpandFrontierTask(frontier_));
  }
  AllocatorDependant * ad_;
  int32 shard_;
  Frontier frontier_;
};

// Insert reachable objects for a shard of an AllocatorDependant
// into a frontier and push it onto the global allocator
struct AboutToCollectTask : public JobQueue::Task {
  AboutToCollectTask (AllocatorDependant * ad) : ad_(ad) {}
  void Run() { ad_->AboutToCollect(); }
  AllocatorDependant * ad_;
};

void GlobalAllocator::AnnounceAboutToCollect() { 
  // Create the initial frontiers
  JobQueue::Job * about_to_collect_job = new JobQueue::Job(G_JQ);
  forall (run, dependants_) {
    AllocatorDependant * ad = *run;
    about_to_collect_job->AddTask(new AboutToCollectTask(ad));
  }
  about_to_collect_job->Wait();
  delete about_to_collect_job;
}

void GlobalAllocator::ComputeReachable() {
  cout << "::ComputeReachable" << endl;
  memory_used_ = 0; boxes_used_ = 0;

  // Create the initial frontiers
  propagate_job_ = new JobQueue::Job(G_JQ);
  forall (run, dependants_) {
    AllocatorDependant * ad = *run;
    int32 num_shards = ad->NumShards();
    for (int32 c=0; c<num_shards; c++)
      propagate_job_->AddTask(new FlagReferencesTask(ad, c));
  }
  propagate_job_->Wait();
  delete propagate_job_;
}

struct ReindexTask : public BlockTask {
  void Handle(Box * o) { 
    o->Index(); 
  }
};

void GlobalAllocator::ReIndex(bool clear) {
  uint64 t0 = TimeNS();
  if (clear) ClearBoxIndex();
  RipAllBoxes<ReindexTask>();
  uint64 t1 = TimeNS();
  cout << "t: " << 1e-9 * (t1 - t0) << " ReIndex::" << endl;
}

void GlobalAllocator::Store(string fn) {
  // GarbageCollect(KILL_THREADS_URGENCY);
  File::Write(fn, block_->StoredMemblock());
}

void GlobalAllocator::Load(string fn) {
  uint64 size = File::Size(fn);
  uint64 capacity = block_->capacity_;
  File::Read(fn, Memblock((char *)block_, size));
  block_->capacity_ = capacity;
  ReIndex(false);
}

struct ZeroReachabilityTask : public BlockTask {
  void Handle(Box * o) { 
    o->RemoveFlag(FLAG_KEEP); 
    o->link_ = NULL;
  } 
};

// Move all boxes in a block towards the beginning of the block
// over the unreachable boxes.
struct CompressBlockTask : public BlockTask {
  void Handle(Box *b) {
    if (!(b->GetFlag(FLAG_KEEP) || (b->GetType() == SYMBOL)))
      return;
    b->RemoveFlag(FLAG_KEEP);
    b->link_ = NULL;
    uint64 size = b->WordSize();
    CHECK(next_ <= (uint64 *)b);
    memmove(next_, b, sizeof(uint64) * size);
    next_ += size; }

  void StartBlock() { next_ = block_->Data(); }
  void FinishBlock() { block_->used_ = next_ - block_->Data(); }
  uint64 * next_;
};

void GlobalAllocator::GarbageCollect(uint64 initial_urgency) {
  cout << "::GarbageCollect" << endl;
  urgency_ = initial_urgency;
  // RipAllBoxes<ZeroReachabilityTask>();

  // Escalate urgency to free up more memory
  while (1) {
    if (urgency_ == KILL_PROCESS_URGENCY) {
      cout << "::GarbageCollect KILL_PROCESS_URGENCY" << endl;
      CHECK(false);
    }
    cout << "urgency_: " << urgency_ << endl;
    AnnounceAboutToCollect();
    ComputeReachable();
    if (HasEnoughMemory()) break;
    urgency_++;
    RipAllBoxes<ZeroReachabilityTask>();
  }

  RipAllBoxes<CompressBlockTask>();
  MoveAllocationBlocks();
  ReIndex();

  cout << "GarbageCollect:: " << endl;
  needs_garbage_collection_ = false;
}

void GlobalAllocator::MoveAllocationBlocks() {
  // b iterates over the second level blocks
  AllocationBlock * b = (AllocationBlock *) block_->Data();
  AllocationBlock * target = b;
  uint64 target_used = 0;
  while (b < (AllocationBlock *) block_->DataEnd()) {
    AllocationBlock * next = b->Next();
    uint64 size = b->used_;
    uint64 * data = b->Data();
    if (target_used >= block_size_) {
      // start new block
      target = target->Next();
      target_used = 0;
    }
    JobQueueMove(G_JQ, target->Data() + target_used, data, size,
		 (uint64 *)G_BOX_INDEX, G_BOX_INDEX_SPACE);
    target_used += size;
    target->capacity_ = target->used_ = target_used;
    b = next;
  }
  // CHECK(target_used > 0); WHAT WAS THIS CHECK FOR? 
  target = target->Next();  
  block_->used_ = (uint64 *)target - block_->Data();
  cout << "Moved blocks - used = " << block_->used_ << endl;
}

bool GlobalAllocator::HasEnoughMemory() {
  cout << "memory_used_:" << memory_used_ << endl;
  cout << "boxes_used_:" << boxes_used_ << endl;
  return ((memory_used_ * 2 < box_space_) &&
	  (boxes_used_ * 4 < G_BOX_INDEX_SPACE));
}

uint64 GlobalAllocator::AccurateRunningMemoryUsed() {
  // This is approximate
  uint64 ret = block_->used_;
  for (AllocationBlock * b = (AllocationBlock *)(G_ALLOC->block_->Data()); 
       b < (AllocationBlock *)(G_ALLOC->block_->DataEnd()); 
       b = b->Next())
    ret -= (b->capacity_ - b->used_);
  return ret;
}

void * GarbageCollectMaintenance(void * junk) {
  if (G_ALLOC->needs_garbage_collection_)
    G_ALLOC->GarbageCollect();  
  return NULL; }

void Frontier::Add(Thing o) {
  if (!o.data_ || !o.IsBoxed()) return;
  Box * b = o.GetBox(); 
  if (!b) { cout << "!b\n"; CHECK(false); }
  if (b->IsMutable()) { cout << "b mutable\n"; CHECK(false); }

  // If it was previously set we're done
  if (b->AtomicSetFlag(FLAG_KEEP)) return;

  seen_boxes_++;
  seen_memory_ += b->WordSize();
  if (b->GetType() != IMMUTABLE_ARRAY) return;
  AddToQueue(b);
}

// adds box b to the queue if it is not already in another queue. 
void Frontier::AddToQueue(Box *b) {
  // Add Arrays to the queue 
  if (__sync_bool_compare_and_swap(&(b->link_), NULL, head_)) {
    head_ = b;
    list_expansion_size_ += b->ArraySize();
  }
}

// Add the range of Objects into the linked list
void Frontier::ExpandRange(const Thing * begin, const Thing * end) {
  for (const Thing * run = begin; run != end; run++) 
    Add(*run);
}

void Frontier::Expand() {
  if (head_ == (Box *)0x1) return;
  Box * head = head_;
  head_ = head_->link_;
  ExpandRange(head->ImmutableArrayData(), 
	      head->ImmutableArrayData() + head->ArraySize());
}

